
import os
import random
import shutil

# 设置原始数据集路径
image_folder = '/home/hw/dataset/cone_total_new/images'
label_folder = '/home/hw/dataset/cone_total_new/labels'
output_path = '/home/hw/dataset/cone'

# 创建输出目录
os.makedirs(os.path.join(output_path, 'train', 'images'), exist_ok=True)
os.makedirs(os.path.join(output_path, 'train', 'labels'), exist_ok=True)
os.makedirs(os.path.join(output_path, 'val', 'images'), exist_ok=True)
os.makedirs(os.path.join(output_path, 'val', 'labels'), exist_ok=True)
os.makedirs(os.path.join(output_path, 'test', 'images'), exist_ok=True)
os.makedirs(os.path.join(output_path, 'test', 'labels'), exist_ok=True)

# 获取所有图像文件
image_files = [f for f in os.listdir(image_folder) if f.endswith('.jpg') or f.endswith('.png')]
random.shuffle(image_files)

# 计算划分的索引
total_images = len(image_files)
train_end = int(0.7 * total_images)
val_end = train_end + int(0.2 * total_images)

# 划分数据集
train_files = image_files[:train_end]
val_files = image_files[train_end:val_end]
test_files = image_files[val_end:]

# 复制文件到相应的目录
def copy_files(file_list, subfolder):
    for file in file_list:
        # 复制图像
        shutil.copy(os.path.join(image_folder, file), os.path.join(output_path, subfolder, 'images', file))
        print(os.path.join(image_folder, file),"->", os.path.join(output_path, subfolder, 'images', file))
        # 复制标签
        label_file = file.replace('.jpg', '.txt').replace('.png', '.txt')
        shutil.copy(os.path.join(label_folder, label_file), os.path.join(output_path, subfolder, 'labels', label_file))
        print(os.path.join(label_folder, label_file),"->",os.path.join(output_path, subfolder, 'labels', label_file))

copy_files(train_files, 'train')
copy_files(val_files, 'val')
copy_files(test_files, 'test')

print("数据集划分完成！")
